#include "nano_v2.h"

using namespace godot;

PackedByteArray NanoV2::encodeStr(const String& str) {
	PackedByteArray result;
	for(int i=0; i < str.length(); i++) {
		uint16_t charCode = static_cast<uint16_t>(str.utf8().get_data()[i]);
		if(charCode <= 0x7f) {
			result.push_back(static_cast<uint8_t>(charCode));
		} else if(charCode <= 0x7ff) {
			result.push_back(static_cast<uint8_t>(0xc0 | (charCode >> 6)));
			result.push_back(static_cast<uint8_t>(0x80 | (charCode & 0x3f)));
		} else {
			result.push_back(static_cast<uint8_t>(0xe0 | (charCode >> 12)));
			result.push_back(static_cast<uint8_t>(0x80 | ((charCode >> 6) & 0x3f)));
			result.push_back(static_cast<uint8_t>(0x80 | (charCode & 0x3f)));
		}
	}
	return result;
}

String NanoV2::decodeStr(const PackedByteArray& buffer) {
	String result;
	size_t offset = 0;
	while (offset < buffer.size()) {
		uint32_t charCode;
		if (buffer[offset] < 128) {
			charCode = buffer[offset++];
		}
		else if (buffer[offset] < 224) {
			charCode = ((buffer[offset] & 0x3f) << 6) + (buffer[offset + 1] & 0x3f);
			offset += 2;
		}
		else {
			charCode = ((buffer[offset] & 0x0f) << 12) +
				((buffer[offset + 1] & 0x3f) << 6) +
				(buffer[offset + 2] & 0x3f);
			offset += 3;
		}
		result += static_cast<char>(charCode);
	}
	return result;
}

PackedByteArray NanoV2::encodePackage(const Ref<NanoPackage>& p_pkg) {
	return encodePackageDirect(p_pkg->type, p_pkg->body);
}

PackedByteArray NanoV2::encodePackageDirect(NanoPackage::PackageType type, const PackedByteArray& body)
{
	PackedByteArray buffer;
	buffer.resize(PKG_HEAD_BYTES + body.size());
	size_t index = 0;

	buffer[index++] = static_cast<uint8_t>(type & 0xff);
	buffer[index++] = static_cast<uint8_t>((body.size() >> 16) & 0xff);
	buffer[index++] = static_cast<uint8_t>((body.size() >> 8) & 0xff);
	buffer[index++] = static_cast<uint8_t>(body.size() & 0xff);

	memcpy(buffer.ptrw() + index, body.ptr(), body.size());

	return buffer;
}

Array NanoV2::decodePackage(const PackedByteArray& buffer) {
	Array result;
	size_t offset = 0;

	while (offset < buffer.size() - 3) {//min package size is 4 bytes
		Ref<NanoPackage> pkg;
		pkg.instantiate();
		pkg->type = static_cast<NanoPackage::PackageType>(buffer[offset++]);

		size_t body_size = (static_cast<size_t>(buffer[offset++]) << 16) |
			(static_cast<size_t>(buffer[offset++]) << 8) |
			static_cast<size_t>(buffer[offset++]) & 0xFFFFFFFF;


		if (body_size > 0 && offset + body_size <= buffer.size()) {
			pkg->body.resize(body_size);
			memcpy(pkg->body.ptrw(), buffer.ptr() + offset, body_size);
			offset += body_size;
		}
		else {
			//pkg->body = PackedByteArray();//no need
			if (body_size != 0)pkg->valid = false;//add
			offset = buffer.size();
		}
		result.append(pkg);
	}
	return result;
}

PackedByteArray NanoV2::encodeMessage(const Ref<NanoMessage>& p_msg) {
	return encodeMessageDirect(p_msg->type, p_msg->id, p_msg->compress_route, p_msg->route, p_msg->body);
}

PackedByteArray NanoV2::encodeMessageDirect(NanoMessage::MessageType type, uint32_t id, bool compressRoute, 
	const String& route, const PackedByteArray& msg)
{
	int idBytes = msgHasId(type) ? calculateMsgIdBytes(id) : 0;
	size_t msgLen = MSG_FLAG_BYTES + idBytes;

	if (msgHasRoute(type)) {
		if (compressRoute) {
			msgLen += MSG_ROUTE_CODE_BYTES;
		}
		else {
			msgLen += MSG_ROUTE_LEN_BYTES + route.length();
		}
	}

	if (!msg.is_empty()) {
		msgLen += msg.size();
	}

	PackedByteArray buffer;
	buffer.resize(msgLen);
	size_t offset = 0;

	offset = encodeMsgFlag(type, compressRoute, buffer, offset);

	if (msgHasId(type)) {
		offset = encodeMsgId(id, buffer, offset);
	}

	if (msgHasRoute(type)) {
		offset = encodeMsgRoute(compressRoute, route.utf8().get_data(), buffer, offset);
	}

	if (!msg.is_empty()) {
		offset = encodeMsgBody(msg, buffer, offset);
	}

	return buffer;
}

Ref<NanoMessage> NanoV2::decodeMessage(const PackedByteArray& buffer) {
	Ref<NanoMessage> msg;
	msg.instantiate();
	size_t offset = 0;
	
	if (buffer.is_empty()) {
		return msg;
	}

	uint8_t flag = buffer[offset++];
	msg->compress_route = flag & MSG_COMPRESS_ROUTE_MASK;
	msg->type = static_cast<NanoMessage::MessageType>((flag >> 1) & MSG_TYPE_MASK);


	if (msgHasId(msg->type)) {
		if (buffer.size() < 2) { msg->valid = false; return msg; }//add
		msg->id = 0;
		int i = 0;
		uint8_t m;
		do {
			m = buffer[offset++];
			msg->id += (static_cast<uint32_t>(m & 0x7f) << (7 * i));
			i++;
			if ((m >= 128 && offset >= buffer.size())||i>NanoV2::MSG_ID_MAX_BYTES) { msg->valid = false; return msg; }//add
		} while (m >= 128 && offset < buffer.size());
	}
	else {
		msg->id = 0;
	}

	if (msgHasRoute(msg->type)) {
		if (msg->compress_route) {

			if (offset + 1 < buffer.size()) {
				uint16_t routeCode = (static_cast<uint16_t>(buffer[offset++]) << 8) | buffer[offset++];
				msg->route = std::to_string(routeCode).c_str();
			}
			else { msg->valid = false; return msg; }//add
		}
		else {
			if (offset < buffer.size()) {
				uint8_t routeLen = buffer[offset++];
				if (routeLen > 0 && offset + routeLen <= buffer.size()) {
					//ByteArray routeBytes(buffer.begin() + offset, buffer.begin() + offset + routeLen);
					PackedByteArray routeBytes;
					routeBytes.resize(routeLen);
					memcpy(routeBytes.ptrw(), buffer.ptr() + offset, routeLen);
					msg->route = decodeStr(routeBytes);
					offset += routeLen;
				}
				else { msg->valid = false; return msg; }//add
			}
			else { msg->valid = false; return msg; }//add
		}

	}

	if (offset < buffer.size()) {
		msg->body.resize(buffer.size() - offset);
		memcpy(msg->body.ptrw(), buffer.ptr() + offset, buffer.size() - offset);
	}

	return msg;
}

void NanoV2::_bind_methods() {
	ClassDB::bind_static_method("NanoV2", D_METHOD("encodeStr", "str"), &NanoV2::encodeStr);
	ClassDB::bind_static_method("NanoV2", D_METHOD("decodeStr", "buffer"), &NanoV2::decodeStr);
	ClassDB::bind_static_method("NanoV2", D_METHOD("encodePackage", "p_pkg"), &NanoV2::encodePackage);
	ClassDB::bind_static_method("NanoV2", D_METHOD("encodePackageDirect", "type", "body"), &NanoV2::encodePackageDirect);
	ClassDB::bind_static_method("NanoV2", D_METHOD("decodePackage", "buffer"), &NanoV2::decodePackage);
	ClassDB::bind_static_method("NanoV2", D_METHOD("encodeMessage", "p_msg"), &NanoV2::encodeMessage);
	ClassDB::bind_static_method("NanoV2", D_METHOD("encodeMessageDirect", "type", "id", "compressRoute", "route", "msg"), &NanoV2::encodeMessageDirect);
	ClassDB::bind_static_method("NanoV2", D_METHOD("decodeMessage", "buffer"), &NanoV2::decodeMessage);
	ClassDB::bind_static_method("NanoV2", D_METHOD("judgeMessageValid", "p_msg"), &NanoV2::judgeMessageValid);
	ClassDB::bind_static_method("NanoV2", D_METHOD("judgePackageValid", "p_pkg"), &NanoV2::judgePackageValid);

}

bool NanoV2::msgHasId(NanoMessage::MessageType type) {
	return type == NanoMessage::TYPE_REQUEST || type == NanoMessage::TYPE_RESPONSE;
}

bool NanoV2::msgHasRoute(NanoMessage::MessageType type) {
	return type == NanoMessage::TYPE_REQUEST || type == NanoMessage::TYPE_NOTIFY || type == NanoMessage::TYPE_PUSH;
}

int NanoV2::calculateMsgIdBytes(uint32_t id) {
	int len = 0;
	do {
		len++;
		id >>= 7;
	} while (id > 0);
	return len;
}

size_t NanoV2::encodeMsgFlag(NanoMessage::MessageType type, bool compressRoute, PackedByteArray& buffer, size_t offset) {
	buffer[offset] = static_cast<uint8_t>((type << 1) | (compressRoute ? 1 : 0));
	return offset + MSG_FLAG_BYTES;
}

size_t NanoV2::encodeMsgId(uint32_t id, PackedByteArray& buffer, size_t offset) {
	do {
		uint8_t tmp = static_cast<uint8_t>(id % 128);
		uint32_t next = id / 128;

		if (next != 0) {
			tmp |= 0x80;
		}

		buffer[offset++] = tmp;
		id = next;
	} while (id != 0);

	return offset;
}

size_t NanoV2::encodeMsgRoute(bool compressRoute, const String& route, PackedByteArray& buffer, size_t offset) {
	if (compressRoute) {

		uint16_t routeCode = 0;//!!
		buffer[offset++] = static_cast<uint8_t>((routeCode >> 8) & 0xff);
		buffer[offset++] = static_cast<uint8_t>(routeCode & 0xff);
	}
	else {
		PackedByteArray routeBytes = encodeStr(route);
		if (routeBytes.size() > 255) {
			throw runtime_error("Route length exceeds maximum limit");
		}

		buffer[offset++] = static_cast<uint8_t>(routeBytes.size() & 0xff);
		//copy(routeBytes.begin(), routeBytes.end(), buffer.begin() + offset);
		memcpy(buffer.ptrw() + offset, routeBytes.ptrw(), routeBytes.size());
		offset += routeBytes.size();
	}

	return offset;
}

size_t NanoV2::encodeMsgBody(const PackedByteArray& msg, PackedByteArray& buffer, size_t offset) {
	//copy(msg.begin(), msg.end(), buffer.begin() + offset);
	memcpy(buffer.ptrw() + offset, msg.ptr(), msg.size());
	return offset + msg.size();
}